from copy import deepcopy
from queue import PriorityQueue, Queue
from typing import Callable, Dict, Hashable, Iterator, List, Optional, Tuple, Union

from .graph import Graph


class SupportedAlgorithms:
    BFS = "bfs"
    DIJKSTRA = "dijkstra"
    A_START = "a_star"


class PathNotExistError(Exception):
    pass


class PathSearchAlg:
    def __init__(self) -> None:
        self._alg = None
        self._come_from = None

    def register(self, alg: str):
        if alg == "bfs":
            self._alg = self._breadth_first_search
        elif alg == "dijkstra":
            self._alg = self._dijkstra_search
        elif alg == "a_star":
            self._alg = self._a_star_search
        else:
            raise NotImplementedError

    @property
    def come_from(self) -> Dict[Hashable, Hashable]:
        return deepcopy(self._come_from)

    def search(
        self,
        graph: Graph,
        start: Hashable,
        end: Optional[Hashable] = None,
        heuristic: Callable[[Hashable, Hashable], int] = None,
    ) -> Tuple[List[Hashable], int, dict[Hashable, Hashable]]:
        self._start = start
        self._come_from = self._alg(graph, start, end, heuristic)
        return self.come_from

    def construct_path(
        self, end: Hashable, raise_error: bool = True
    ) -> Union[List[Hashable], None]:
        assert self._start is not None, "Search algorithm not run"
        if end not in self._come_from:
            if raise_error:
                raise PathNotExistError()
            else:
                return None
        path = []
        current = end
        while current != self._start:
            path.append(current)
            current = self._come_from[current]
        path.append(self._start)
        path.reverse()
        return path

    def _breadth_first_search(
        self, graph: Graph, start: Hashable, end: Optional[Hashable], *args, **kwargs
    ) -> Tuple[List[Hashable], int, dict[Hashable, Hashable]]:
        come_from = {start: None}
        frontier = Queue()
        frontier.put(start)

        while not frontier.empty():
            current = frontier.get()
            if end and current == end:
                break
            for node in graph.neighbors(current):
                if node not in come_from.keys():
                    frontier.put(node)
                    come_from[node] = current
        return come_from

    def _dijkstra_search(
        self, graph: Graph, start: Hashable, end: Optional[Hashable], *args, **kwargs
    ) -> Tuple[List[Hashable], int, dict[Hashable, Hashable]]:
        frontier = PriorityQueue()
        frontier.put((0, start))
        come_from = {start: None}
        cost_so_far = {start: 0}

        while not frontier.empty():
            _, current = frontier.get()
            if end and current == end:
                break
            for node in graph.neighbors(current):
                new_cost = cost_so_far[current] + graph.get_edge_data(current, node)
                if node not in cost_so_far or new_cost < cost_so_far[node]:
                    cost_so_far[node] = new_cost
                    priority = new_cost
                    frontier.put((priority, node))
                    come_from[node] = current
        return come_from

    def _a_star_search(
        self,
        graph: Graph,
        start: Hashable,
        end: Hashable,
        heuristic: Callable[[Hashable, Hashable], float],
        *args,
        **kwargs,
    ) -> Tuple[List[Hashable], int, dict[Hashable, Hashable]]:
        frontier = PriorityQueue()
        frontier.put((0, start))
        come_from = {start: None}
        cost_so_far = {start: 0}

        while not frontier.empty():
            _, current = frontier.get()
            if current == end:
                break
            for node in graph.neighbors(current):
                new_cost = cost_so_far[current] + heuristic(end, node)
                if node not in cost_so_far or new_cost < cost_so_far[node]:
                    cost_so_far[node] = new_cost
                    priority = new_cost + heuristic(end, node)
                    frontier.put((priority, node))
                    come_from[node] = current

        return come_from
